/*
 *
 * Copyright 2019 Asylo authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

#ifndef ASYLO_IDENTITY_PROVISIONING_SGX_INTERNAL_FAKE_SGX_PCS_CLIENT_H_
#define ASYLO_IDENTITY_PROVISIONING_SGX_INTERNAL_FAKE_SGX_PCS_CLIENT_H_

#include <atomic>
#include <memory>
#include <string>
#include <vector>

#include "absl/container/flat_hash_map.h"
#include "absl/hash/hash.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "asylo/crypto/certificate.pb.h"
#include "asylo/crypto/signing_key.h"
#include "asylo/crypto/x509_certificate.h"
#include "asylo/identity/platform/sgx/machine_configuration.pb.h"
#include "asylo/identity/provisioning/sgx/internal/container_util.h"
#include "asylo/identity/provisioning/sgx/internal/pck_certificate_util.h"
#include "asylo/identity/provisioning/sgx/internal/platform_provisioning.pb.h"
#include "asylo/identity/provisioning/sgx/internal/sgx_pcs_client.h"
#include "asylo/identity/provisioning/sgx/internal/sgx_pcs_client.pb.h"
#include "asylo/identity/provisioning/sgx/internal/tcb.pb.h"
#include "asylo/util/mutex_guarded.h"
#include "asylo/util/statusor.h"

namespace asylo {
namespace sgx {

// A fake implementation of SgxPcsClient. FakeSgxPcsClient is thread-safe. Each
// FakeSgxPcsClient stores a map from a set of known FMSPCs to TCB infos.
//
// The FakeSgxPcsClient implementation uses fake layouts for SGX identifiers to
// embed information about the platforms they represent. This information is
// needed to implement the methods of SgxPcsClient without the use of a database
// of SGX identifiers. Fake SGX identifiers generated by one FakeSgxPcsClient
// can be understood by another FakeSgxPcsClient. The FakeSgxPcsClients need not
// be in the same process or on the same machine.
//
// Fake data generated by FakeSgxPcsClient obeys all the constraints documented
// in the specifications for the corresponding data structures defined in the
// SGX specifications, except that the fake certificates and TCB info use
// different keys and CA certificates for the SGX PKI. The structure of the
// hierarchy of the fake PKI matches the hierarchy of real SGX PKI.
//
// FakeSgxPcsClient relies on an internal fake layout for FMSPCs and PPIDs.
class FakeSgxPcsClient : public SgxPcsClient {
 public:
  // A fake SGX platform type. The information in this struct is embedded in
  // each fake FMSPC created for this type.
  struct PlatformProperties {
    // The CA that issues this platform's PCK certificates.
    SgxCaType ca;

    // The PCE ID used by the platform.
    PceId pce_id;
  };

  // Constructs a new FakeSgxPcsClient using the fake PKI and PCK in
  // fake_sgx_pki.h.
  FakeSgxPcsClient();

  // Constructs a new FakeSgxPcsClient using the fake PKI in fake_sgx_pki.h and
  // |pck_der| as the public PCK in DER form.
  explicit FakeSgxPcsClient(absl::string_view pck_der);

  // Adds |fmspc| to the set of known FMSPCs and associates it with |tcb_info|,
  // which must have a |version| of 2.
  //
  // If the |fmspc| is already in the set of known FMSPCs, then AddFmspc()
  // returns false and does not modify existing data. Otherwise, AddFmspc()
  // returns true.
  //
  // AddFmspc() returns an error if |fmspc| does not match the internal FMSPC
  // layout. To create a matching FMSPC, use CreateFmspcWithProperties().
  StatusOr<bool> AddFmspc(Fmspc fmspc, TcbInfo tcb_info);

  // Changes the TCB info associated with |fmspc| to be |tcb_info|, which must
  // have a |version| of 2.
  //
  // Returns an error if |fmspc| is not valid or is not in the set of known
  // FMSPCs.
  Status UpdateFmspc(const Fmspc &fmspc, TcbInfo tcb_info);

  // From SgxPcsClient. Currently unimplemented.
  StatusOr<GetPckCertificateResult> GetPckCertificate(
      const Ppid &ppid, const CpuSvn &cpu_svn, const PceSvn &pce_svn,
      const PceId &pce_id) override;

  // From SgxPcsClient. Returns a set of PCK certificates for |ppid| and
  // |pce_id| based on the configured TCB info for |ppid|'s platform.
  // Specifically, GetPckCertificates() returns one PCK certificate for each TCB
  // level in the TCB info configured for |ppid|'s platform.
  //
  // The returned certificates follow the SGX PCK certificate specification at
  // https://download.01.org/intel-sgx/dcap-1.0/docs/SGX_PCK_Certificate_CRL_Spec-1.0.pdf.
  //
  // The PCK certificates are re-generated after every call, so two back-to-back
  // calls with the same arguments may return certificates with different
  // signatures.
  //
  // GetPckCertificates() returns an error if:
  //
  //   * |ppid| or |pce_id| are invalid.
  //   * |ppid| does not conform to the internal layout.
  //   * The FMSPC for |ppid| is not in the set of known FMSPCs.
  //   * |ppid| is for an unsupported SgxCaType.
  StatusOr<GetPckCertificatesResult> GetPckCertificates(
      const Ppid &ppid, const PceId &pce_id) override;

  // From SgxPcsClient. Currently unimplemented.
  StatusOr<GetCrlResult> GetCrl(SgxCaType sgx_ca_type) override;

  // From SgxPcsClient. Returns the TCB info associated with |fmspc|, but
  // modified to have an "issueDate" corresponding to the current time and a
  // "nextUpdate" one month later.
  //
  // The returned signatures over the TCB info may differ between calls, but
  // they will all be correct signatures.
  //
  // GetTcbInfo() returns an error if:
  //
  //   * |fmspc| is invalid or does not conform to the internal layout.
  //   * |fmspc| is not in the set of known FMSPCs.
  //
  StatusOr<GetTcbInfoResult> GetTcbInfo(const Fmspc &fmspc) override;

  // Returns a new random FMSPC that a FakeSgxPcsClient can associate with
  // |properties|. Callers should be careful to check for collisions with
  // previous randomly-generated FMSPCs.
  StatusOr<Fmspc> CreateFmspcWithProperties(
      const PlatformProperties &properties);

  // Returns a new random PPID that a FakeSgxPcsClient can associate with
  // |fmspc|.
  static StatusOr<Ppid> CreatePpidForFmspc(const Fmspc &fmspc);

 private:
  // Contains information about a CA in the fake SGX PKI.
  struct CaInfo {
    // The certificate chain beginning with the CA's certificate and ending in
    // the fake SGX root CA certificate. All certificates must be in PEM format.
    CertificateChain issuer_chain;

    // The CA's certificate interface.
    std::unique_ptr<const X509Certificate> certificate;

    // The CA's signing key.
    std::unique_ptr<const SigningKey> signing_key;

    CaInfo() = default;

    // Constructs a CaInfo for the CA given by the chain of PEM-encoded
    // certificates in |issuer_chain_pem| and the PEM-encoded signing key in
    // |signing_key_pem|.
    CaInfo(absl::Span<const absl::string_view> issuer_chain_pem,
           absl::string_view signing_key_pem);
  };

  using FmspcToTcbInfoMap =
      absl::flat_hash_map<Fmspc, TcbInfo, absl::Hash<Fmspc>, MessageEqual>;

  // Returns the CA info for the given |ca_type|, or an INVALID_ARGUMENT error
  // if |ca_type| is not supported. Never returns nullptr.
  static StatusOr<const CaInfo *> GetCaInfo(SgxCaType ca_type);

  // Creates a PckCertificateInfo message for the given |ca_info| and with the
  // given |extension_data|. All the |extension_data| must be valid and
  // consistent, and it all must match the internal provisioning type layouts.
  StatusOr<PckCertificates::PckCertificateInfo> CreateFakePckCertificateInfo(
      const CaInfo &ca_info, const SgxExtensions &extension_data) const;

  // Creates a fake PCK certificate in PEM format for the given |ca_info| and
  // with the given |extension_data|. All the |extension_data| must be valid and
  // consistent, and it all must match the internal provisioning type layouts.
  StatusOr<Certificate> CreateFakePckCertificate(
      const CaInfo &ca_info, const SgxExtensions &extension_data) const;

  const CertificateChain tcb_info_issuer_chain_;
  const std::unique_ptr<const SigningKey> tcb_info_signing_key_;
  const std::string pck_public_key_der_;
  MutexGuarded<FmspcToTcbInfoMap> tcb_infos_;
  std::atomic<uint32_t> fmspc_id_counter_;
};

}  // namespace sgx
}  // namespace asylo

#endif  // ASYLO_IDENTITY_PROVISIONING_SGX_INTERNAL_FAKE_SGX_PCS_CLIENT_H_
